package datastore

import (
	"context"
	"testing"

	namespaceDS "github.com/stackrox/rox/central/namespace/datastore"
	"github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/cache"
	"github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/common"
	"github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/datastore/internal/store"
	pgStore "github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/datastore/internal/store/postgres"
	v1 "github.com/stackrox/rox/generated/api/v1"
	"github.com/stackrox/rox/generated/storage"
	"github.com/stackrox/rox/pkg/postgres"
	"github.com/stackrox/rox/pkg/search"
	"github.com/stackrox/rox/pkg/testutils"
)

// DataStore is an intermediary to VulnerabilityRequest storage.
//
//go:generate mockgen-wrapper
type DataStore interface {
	Search(ctx context.Context, q *v1.Query) ([]search.Result, error)
	SearchRequests(ctx context.Context, q *v1.Query) ([]*v1.SearchResult, error)
	SearchRawRequests(ctx context.Context, q *v1.Query) ([]*storage.VulnerabilityRequest, error)

	Count(ctx context.Context, q *v1.Query) (int, error)
	Exists(ctx context.Context, id string) (bool, error)
	Get(ctx context.Context, id string) (*storage.VulnerabilityRequest, bool, error)
	GetMany(ctx context.Context, ids []string) ([]*storage.VulnerabilityRequest, error)

	// AddRequest adds a new request to the data store.
	AddRequest(ctx context.Context, request *storage.VulnerabilityRequest) error

	// UpdateRequestStatus updates an existing request from pending state to approved or denied.
	UpdateRequestStatus(ctx context.Context, id string, comment string, status storage.RequestStatus) (*storage.VulnerabilityRequest, error)

	// UpdateRequestExpiry updates an existing deferral request's expiry. Used in v1 only.
	// The client must meet one of the following criteria:
	// 1. VulnerabilityManagementApprovals write permission.
	// 2. VulnerabilityManagementRequests write permission and is the requester of the original exception.
	//
	// Use only in v1 workflow
	UpdateRequestExpiry(ctx context.Context, id, message string, updatedExpiry *storage.RequestExpiry) (*storage.VulnerabilityRequest, error)

	// UpdateRequest updates an existing deferral or false positive exception.
	// The client must meet one of the following criteria:
	// 1. VulnerabilityManagementApprovals write permission.
	// 2. VulnerabilityManagementRequests write permission and is the requester of the original exception.
	// INVARIANT: `id` and `reqParams` must be valid.
	//
	// Use only in v2 workflow.
	UpdateRequest(ctx context.Context, id string, update *common.UpdateRequest) (*storage.VulnerabilityRequest, error)

	// MarkRequestInactive marks a vulnerability exception inactive. Once inactive, the exception enforcement is reverted.
	// The client must meet one of the following criteria:
	// 1. VulnerabilityManagementApprovals write permission.
	// 2. VulnerabilityManagementRequests write permission and is the requester of the original exception.
	MarkRequestInactive(ctx context.Context, id string, comment string) (*storage.VulnerabilityRequest, error)

	// RemoveRequest removes the request with specified ID from the database. Only pending exceptions and pending updates
	// to an enforced exception can be deleted through API. All other types of exceptions can only be cancelled.
	// Exceptions are retained in the system according to the retention configuration.
	RemoveRequest(ctx context.Context, id string) error

	//// FOR INTERNAL USE ONLY.

	// RemoveRequestsInternal removes the vulnerability request with specified ID without performing any validation.
	RemoveRequestsInternal(ctx context.Context, ids []string) error
}

// New returns a new instance of DataStore using the input store, and searcher.
func New(s store.Store, pendingReqCache cache.VulnReqCache,
	activeReqCache cache.VulnReqCache, namespaceDatastore namespaceDS.DataStore) DataStore {
	d := &datastoreImpl{
		store:              s,
		pendingReqCache:    pendingReqCache,
		activeReqCache:     activeReqCache,
		namespaceDatastore: namespaceDatastore,
	}
	return d
}

// GetTestPostgresDataStore provides a datastore connected to postgres for testing purposes.
func GetTestPostgresDataStore(t testing.TB, pool postgres.DB, pendingReqCache cache.VulnReqCache, activeReqCache cache.VulnReqCache) (DataStore, error) {
	testutils.MustBeInTest(t)
	namespaceStore, err := namespaceDS.GetTestPostgresDataStore(t, pool)
	if err != nil {
		return nil, err
	}

	storage := pgStore.New(pool)
	d := &datastoreImpl{
		store:              storage,
		pendingReqCache:    pendingReqCache,
		activeReqCache:     activeReqCache,
		namespaceDatastore: namespaceStore,
	}

	return d, nil
}
